
import json
import os
import sqlite3
import openai

import requests

openai.api_key = os.environ["OPENAI_API_KEY_ZHIHU"]
openai.api_base= os.environ["OPENAI_API_BASE_ZHIHU"]


# Example dummy function hard coded to return the same weather
# In production, this could be your backend API or an external API
amap_key = "6d672e6194caa3b639fccf2caf06c342"
cursor = None

database_schema_string = """
CREATE TABLE orders (
    id INT PRIMARY KEY NOT NULL, -- 主键，不允许为空
    customer_id INT NOT NULL, -- 客户ID，不允许为空
    product_id STR NOT NULL, -- 产品ID，不允许为空
    price DECIMAL(10,2) NOT NULL, -- 价格，不允许为空
    status INT NOT NULL, -- 订单状态，整数类型，不允许为空。0代表待支付，1代表已支付，2代表已退款
    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, -- 创建时间，默认为当前时间
    pay_time TIMESTAMP -- 支付时间，可以为空
);
"""
tools = [{
    # 摘自 OpenAI 官方示例 https://github.com/openai/openai-cookbook/blob/main/examples/How_to_call_functions_with_chat_models.ipynb
    "type": "function",
    "function": {
        "name": "ask_database",
        "description": "Use this function to answer user questions about business. \
                    Output should be a fully formed SQL query.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": f"""
                    SQL query extracting info to answer the user's question.
                    SQL should be written using this database schema:
                    {database_schema_string}
                    The query should be returned in plain text, not in JSON.
                    The query should only contain grammars supported by SQLite.
                    """,
                }
            },
            "required": ["query"],
        }
    }
}]

def initDB():
    global cursor
    # 创建数据库连接
    conn = sqlite3.connect(':memory:')
    cursor = conn.cursor()

    # 创建orders表
    cursor.execute(database_schema_string)

    # 插入5条明确的模拟记录
    mock_data = [
        (1, 1001, 'TSHIRT_1', 50.00, 0, '2023-09-12 10:00:00', None),
        (2, 1001, 'TSHIRT_2', 75.50, 1, '2023-09-16 11:00:00', '2023-08-16 12:00:00'),
        (3, 1002, 'SHOES_X2', 25.25, 2, '2023-10-17 12:30:00', '2023-08-17 13:00:00'),
        (4, 1003, 'SHOES_X2', 25.25, 1, '2023-10-17 12:30:00', '2023-08-17 13:00:00'),
        (5, 1003, 'HAT_Z112', 60.75, 1, '2023-10-20 14:00:00', '2023-08-20 15:00:00'),
        (6, 1002, 'WATCH_X001', 90.00, 0, '2023-10-28 16:00:00', None)
    ]

    for record in mock_data:
        cursor.execute('''
        INSERT INTO orders (id, customer_id, product_id, price, status, create_time, pay_time)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        ''', record)

    # 提交事务
    conn.commit()


def ask_database(query):
    global cursor
    cursor.execute(query)
    records = cursor.fetchall()
    return records

def getChatResponse(messages,model="gpt-3.5-turbo"):
    # Step 1: send the conversation and available functions to the model
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0,
        seed=1024,  # 随机种子保持不变，temperature 和 prompt 不变的情况下，输出就会不变
        tool_choice="auto",  # 默认值，由 GPT 自主决定返回 function call 还是返回文字回复。也可以强制要求必须调用指定的函数，详见官方文档
        tools=tools,
    )
    return response

def run_conversation(userInput,model="gpt-3.5-turbo"):
    messages = [{"role": "system", "content": "你是一个数据分析师，基于数据库的数据回答问题"},
                {"role": "system", "content": userInput}]

    response = getChatResponse(messages,model)
    response_message = response.choices[0].message
    tool_calls = response_message.tool_calls
    result_response = response
    # Step 2: check if the model wanted to call a function
    while tool_calls:
        # Step 3: call the function
        # Note: the JSON response may not always be valid; be sure to handle errors
        available_functions = {
            "ask_database": ask_database
        }  # only one function in this example, but you can have multiple
        messages.append(result_response.choices[0].message)  # extend conversation with assistant's reply
        # Step 4: send the info for each function call and function response to the model
        for tool_call in tool_calls:
            function_name = tool_call.function.name
            function_to_call = available_functions[function_name]
            function_args = json.loads(tool_call.function.arguments)
            function_response = function_to_call(**function_args)
            messages.append(
                {
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": str(function_response),
                }
            )  # extend conversation with function response
        second_response = getChatResponse(messages, model)
        #判断是否 second_response.choices[0].message 有tool_calls属性

        if "tool_calls" in second_response.choices[0].message:
            tool_calls = second_response.choices[0].message.tool_calls
        else:
            tool_calls = None

        result_response = second_response

    return result_response

initDB()
result_response = run_conversation("10月的销售额")
print(result_response.choices[0].message.content)